package com.suemi.missile.common.utils;

import com.suemi.missile.common.entity.Patent;
import com.suemi.missile.common.mql.domain.Query;
import com.suemi.missile.common.mql.domain.Where;
import com.suemi.missile.common.mql.exception.MqlParseException;
import com.suemi.missile.common.mql.parse.MqlParseVisitor;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

public class MqlBatchRanker implements PatentRanker {

  private static ExecutorService executor;

  static {
    int cores = Runtime.getRuntime().availableProcessors();
    executor = Executors.newFixedThreadPool(cores);
  }

  private final String mql;

  private final Where where;
  private final List<Patent> patents;
  private List<BatchResult> batches;
  private volatile List<Patent> ranked;
  private volatile List<Patent> discarded;
  private int progressed;
  private boolean stop = false;
  private int batchSize = 400;

  public MqlBatchRanker(String mql, List<Patent> patents) throws MqlParseException {
    this.mql = mql;
    this.patents = patents;
    Query query = MqlParseVisitor.generate(mql);
    where = query.getWhere();
    reset();
  }

  public static void shutdown() {
    executor.shutdown();
  }

  @Override
  public boolean rank() {
    if (!start()) return false;
    int cur = 0;

    Random random = new Random();
    while (cur < patents.size() && !stop) {
      int realSize = batchSize + random.nextInt(batchSize);
      executor.execute(new BatchTask(cur, realSize));
      cur += realSize;
    }
    return true;
  }

  @Override
  public synchronized boolean done() {
    return progressed == patents.size();
  }

  private synchronized boolean start() {
    if (progressed != -1) return false;
    progressed = 0;
    return true;
  }

  @Override
  public List<Patent> accept() {
    if (!done()) return null;
    if (ranked != null) return ranked;

    List<Pair<Patent, Float>> tmp = new ArrayList<>();

    final List<BatchResult> tmpBatches = batches;
    tmpBatches.forEach(b -> tmp.addAll(b.hits));

    tmp.sort(new Comparator<Pair<Patent, Float>>() {
      @Override
      public int compare(Pair<Patent, Float> o1, Pair<Patent, Float> o2) {
        return o2.getRight().compareTo(o1.getRight());
      }
    });

    if (ranked == null) ranked = tmp.stream().map(Pair::getLeft).collect(Collectors.toList());
    return ranked;
  }

  @Override
  public List<Patent> discard() {
    if (!done()) return null;
    if (discarded != null) return discarded;

    List<Patent> tmp = new ArrayList<>();
    final List<BatchResult> tmpBatches = batches;
    tmpBatches.forEach(b -> tmp.addAll(b.exclude));

    if (discarded == null) discarded = tmp;
    return discarded;
  }

  @Override
  public List<Patent> track(int start, int limit) {
    assert start >= 0 && limit >= 0;
    if (done()) {
      if (ranked == null) accept();
      int end = Math.min(start + limit, ranked.size());
      return ranked.subList(start, end);
    }

    int cur = 0;
    List<Patent> res = new ArrayList<>();

    List<BatchResult> tmpBatches;
    synchronized (this) {
      tmpBatches = new ArrayList<>(batches);
    }

    for (BatchResult batch : tmpBatches) {
      if (cur + batch.num <= start) continue;
      int end = Math.min(limit, batch.num + cur - start);
      res.addAll(batch.hits.subList(start - cur, end).stream().map(Pair::getLeft).collect(Collectors.toList()));
      limit -= (end + cur - start);
      if (limit <= 0) break;
    }

    return res;
  }

  @Override
  public synchronized void cancel() {
    stop = true;
  }

  private void reset() {
    progressed = -1;
    batches = new ArrayList<>();
    stop = false;
    ranked = null;
  }

  public synchronized int getProgressed() {
    return progressed;
  }

  private static class BatchResult {
    public int num;

    public List<Pair<Patent, Float>> hits = new ArrayList<>();

    public List<Patent> exclude = new ArrayList<>();
  }

  private class BatchTask implements Runnable {

    public int start;
    public int size;

    public BatchTask(int start, int size) {
      this.start = start;
      this.size = size;
    }

    @Override
    public void run() {
      int end = Math.min(start + size, patents.size());
      BatchResult batch = new BatchResult();
      for (int cur = start; cur < end; ++cur) {
        Patent patent = patents.get(cur);
        float score = where.genScore(patent);
        if (score > 0) {
          ++batch.num;
          batch.hits.add(new ImmutablePair<>(patent, score));
        } else
          batch.exclude.add(patent);
      }

      synchronized (MqlBatchRanker.this) {
        batches.add(batch);
        progressed += (end - start);
      }
    }
  }
}
